{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mf3kOv1YMB5y"
      },
      "source": [
        "# $\\hat R$ locker\n",
        "\n",
        "This notebook serves as a sandbox to understand the potential of the nested-$\\hat R$ diagnostic. The underlying idea is to gather short chains into a long \"super chains\" and then check that the super chains are mixing. We'll motivate this idea and work out details, benefits and limitations. \n",
        "\n",
        "Copyright 2021 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-rOdskBSMfQN"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLYQYt7BNiFj"
      },
      "outputs": [],
      "source": [
        "# Import tf first to enable eager mode.\n",
        "import tensorflow as tf\n",
        "tf.executing_eagerly()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z2hr9nAAC2aB"
      },
      "outputs": [],
      "source": [
        "# TODO (charlesm93): check which of these actually need to be imported.\n",
        "\n",
        "import numpy as np\n",
        "from matplotlib.pyplot import *\n",
        "# %config InlineBackend.figure_format = 'retina'\n",
        "# matplotlib.pyplot.style.use(\"dark_background\")\n",
        "\n",
        "import jax\n",
        "from jax import random\n",
        "from jax import numpy as jnp\n",
        "\n",
        "from colabtools import adhoc_import\n",
        "\n",
        "from inference_gym import using_jax as gym\n",
        "\n",
        "from tensorflow_probability.spinoffs.fun_mc import using_jax as fun_mcmc\n",
        "\n",
        "\n",
        "# import tensorflow as tf\n",
        "from tensorflow_probability.python.internal import prefer_static as ps\n",
        "from tensorflow_probability.python.internal import unnest\n",
        "\n",
        "\n",
        "import tensorflow_probability as _tfp\n",
        "tfp = _tfp.substrates.jax\n",
        "tfd = tfp.distributions\n",
        "tfb = tfp.bijectors\n",
        "\n",
        "tfp_np = _tfp.substrates.numpy\n",
        "tfd_np = tfp_np.distributions "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ru7PwrbIUETu"
      },
      "outputs": [],
      "source": [
        "# set font size for matplot lib\n",
        "font = {'family' : 'normal',\n",
        "        'weight' : 'bold',\n",
        "        'size'   : 14}\n",
        "\n",
        "matplotlib.rc('font', **font)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LkCpkVpuMrdF"
      },
      "outputs": [],
      "source": [
        "tf.executing_eagerly()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "voAzU660FbPd"
      },
      "source": [
        "## Set up problem"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yUtqH620FdYf"
      },
      "outputs": [],
      "source": [
        "# options: Bananas, GermanCredit, Brownian\n",
        "problem_name = 'Bananas'\n",
        "\n",
        "if (problem_name == 'Bananas'):\n",
        "  target = gym.targets.VectorModel(gym.targets.Banana(),\n",
        "                                    flatten_sample_transformations=True)\n",
        "  num_dimensions = target.event_shape[0]  \n",
        "  init_step_size = 1.\n",
        "\n",
        "if (problem_name == 'GermanCredit'):\n",
        "  # This problem seems to require that we load TF datasets first.\n",
        "  import tensorflow_datasets\n",
        "  target = gym.targets.VectorModel(gym.targets.GermanCreditNumericSparseLogisticRegression(),\n",
        "                                    flatten_sample_transformations=True)\n",
        "  num_dimensions = target.event_shape[0]\n",
        "  init_step_size = 0.02\n",
        "\n",
        "if (problem_name == 'Brownian'):\n",
        "  target = gym.targets.BrownianMotionMissingMiddleObservations()\n",
        "  target = gym.targets.VectorModel(target,\n",
        "                                    flatten_sample_transformations = True)\n",
        "  num_dimensions = target.event_shape[0]\n",
        "  init_step_size = 0.01\n",
        "\n",
        "def target_log_prob_fn(x):\n",
        "  \"\"\"Unnormalized, unconstrained target density.\n",
        "\n",
        "  This is a thin wrapper that applies the default bijectors so that we can\n",
        "  ignore any constraints.\n",
        "  \"\"\"\n",
        "  y = target.default_event_space_bijector(x)\n",
        "  fldj = target.default_event_space_bijector.forward_log_det_jacobian(x)\n",
        "  return target.unnormalized_log_prob(y) + fldj\n",
        "\n",
        "# NOTE: use a large factor to get overdispered initializations.\n",
        "# NOTE: don't set offset to 0 when the target mean is 0.\n",
        "# CHECK: what scale should we use? Poor inits can make the problem much more\n",
        "# difficult.\n",
        "# NOTE: we probably want inits that allow us to get decent estimates\n",
        "# in the long regime\n",
        "\n",
        "# if (problem_name == 'Bananas'):\n",
        "if (problem_name == 'Bananas'):\n",
        "  offset = 2\n",
        "  def initialize (shape, key = random.PRNGKey(37272709)):\n",
        "    return 3 * random.normal(key, shape + (num_dimensions,)) + offset\n",
        "\n",
        "if (problem_name == 'GermanCredit'):\n",
        "  offset = 0.1\n",
        "  def initialize (shape, key = random.PRNGKey(37272709)):\n",
        "    return 0.5 * random.normal(key, shape + (num_dimensions,)) + offset\n",
        "\n",
        "# offset = 0.5\n",
        "# def initialize (shape, key = random.PRNGKey(37272709)):\n",
        "#   return 0.01 * random.normal(key, shape + (num_dimensions,)) + offset\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xr5bsemmOvm0"
      },
      "source": [
        "## Run MCMC\n",
        "We consider two regimes: the \"long\" regime in which a few chains are run for many warmup and sampling iterations, and the \"short\" regime, wherein many chains are run for a few warmup and sampling iterations. Note that in the short regime we're willing to not warmup our chains (i.e. possibly adapt step size, trajectory length, mass matrix) as well as in the long regime, the hope being that the variance decreases enough because we're running many chains. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jd5akqq3G6RH"
      },
      "outputs": [],
      "source": [
        "# Transition kernel for long regime\n",
        "num_chains_long = 4\n",
        "if (problem_name == 'GermanCredit'):\n",
        "  num_warmup_long, num_sampling_long = 500, 1000\n",
        "if (problem_name == 'Bananas'):\n",
        "  num_warmup_long, num_sampling_long = 200, 1000\n",
        "total_samples_long = num_warmup_long + num_sampling_long\n",
        "\n",
        "# CHECK: is this the transition kernel we want to use?\n",
        "# REMARK: the step size is picked based on the model we're fitting\n",
        "if (problem_name == 'Bananas' or problem_name == 'GermanCredit'): \n",
        "  kernel_long = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "  kernel_long = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_long, num_warmup_long)\n",
        "  kernel_long = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "      kernel_long, num_warmup_long, target_accept_prob = 0.75,\n",
        "      reduce_fn=tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "# Follow the inference gym tutorial\n",
        "# NOTE: transition kernel below is untested.\n",
        "if (problem_name == 'Brownian'):\n",
        "  kernel_long = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "  # Adapt step size.\n",
        "  kernel_long = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "    kernel_long, num_warmup_long, # int(num_samples // 2 * 0.8),\n",
        "    target_accept_prob = 0.9)\n",
        "  # Adapt trajectory length.\n",
        "  kernel_long = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(\n",
        "    kernel_long,\n",
        "    num_adaptation_steps = num_warmup_long) # int(num_steps // 2 * 0.8))\n",
        "\n",
        "\n",
        "# TODO: work out what an appropriate transition kernel for this problem would be.\n",
        "# if (problem_name == 'GermanCredit'):\n",
        "#   kernel_long = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "#   kernel_long = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_long, num_warmup_long)\n",
        "#   kernel_long = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "#       kernel_long, num_warmup_long, target_accept_prob = 0.75,\n",
        "#       reduce_fn=tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "initial_state = initialize((num_chains_long,))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "elCk0Y_dk_OG"
      },
      "outputs": [],
      "source": [
        "# initial_state = initialize((num_chains_long,))\n",
        "result_long = tfp.mcmc.sample_chain(\n",
        "  total_samples_long, initial_state, kernel = kernel_long, seed = random.PRNGKey(1954))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aZh-AjqiOuN2"
      },
      "outputs": [],
      "source": [
        "# Transition kernel for short regime\n",
        "# CHECK: how many warmup iterations should we use here?\n",
        "# Suggested options: 512, 1024, 2048, 2500\n",
        "num_chains_short = 512\n",
        "num_super_chains = 4\n",
        "\n",
        "if (problem_name == 'GermanCredit'):\n",
        "  num_warmup_short, num_sampling_short = 1000, 1000\n",
        "if (problem_name == 'Bananas'):\n",
        "  num_warmup_short, num_sampling_short = 100, 1000  # 100, 1000\n",
        "total_samples_short = num_warmup_short + num_sampling_short\n",
        "\n",
        "if (problem_name == 'Bananas' or problem_name == 'GermanCredit'):\n",
        "  kernel_short = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "  kernel_short = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_short, num_warmup_short)\n",
        "  kernel_short = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "      kernel_short, num_warmup_short, target_accept_prob = 0.75,  #0.75,\n",
        "      reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "different_location = False\n",
        "\n",
        "if (different_location):\n",
        "  # initialize each chain at a different location \n",
        "  initial_state = initialize((num_chains_short,))\n",
        "\n",
        "else:\n",
        "  # Chains within a super chain are all initialized at the same location\n",
        "  # Here we use the same initial points as in the long regime.\n",
        "  initial_state = initial_state  # initialize((num_super_chains,))\n",
        "  initial_state = np.repeat(initial_state, num_chains_short // num_super_chains,\n",
        "                            axis = 0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WdHDBo15RvJ4"
      },
      "outputs": [],
      "source": [
        "\n",
        "result_short = tfp.mcmc.sample_chain(\n",
        "    total_samples_short, initial_state, kernel = kernel_short,\n",
        "    seed = random.PRNGKey(1954))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B-kT_5eAlEU7"
      },
      "source": [
        "## Analyze results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cCdWa58uGL-w"
      },
      "source": [
        "### Squared error for Monte Carlo estimate of the mean and variance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VXf5CfXSjWSB"
      },
      "outputs": [],
      "source": [
        "# Get some estimates of the mean and variance.\n",
        "try:\n",
        "  mean_est = target.sample_transformations['identity'].ground_truth_mean\n",
        "except:\n",
        "  print('no ground truth mean')\n",
        "  mean_est = (result.all_states[num_warmup:, :]).mean(0).mean(0)\n",
        "try:\n",
        "  var_est = target.sample_transformations['identity'].ground_truth_standard_deviation**2\n",
        "except:\n",
        "  print('no ground truth std dev')\n",
        "  var_est = ((result.all_states[num_warmup:, :]**2).mean(0).mean(0) -\n",
        "             mean_est**2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TASA9Hbq83dv"
      },
      "outputs": [],
      "source": [
        "jnp.linalg.norm(var_est[0] / 100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pD2srL4Cmuoh"
      },
      "source": [
        "As a first step plot the squared error based on a Monte Carlo estimator that discards the first half of the samples, and doesn't discriminate between warmup and sampling iterations. We also plot the target precision whith the \"true\" variance -- when available for instance via the inference gym -- divided by 100. This is the precision we expect our Monte Carlo estimates to reach with an effective sample size of 100."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kuoSPJOEUqdb"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Map MCMC samples from the unconstrained space to the original space\n",
        "# CHECK: does this mess up the banana example?\n",
        "result_state_long = target.default_event_space_bijector(result_long.all_states)\n",
        "result_state_short = target.default_event_space_bijector(result_short.all_states)\n",
        "\n",
        "def mc_est(x, axis = 0):\n",
        "  \"\"\"Computes the running sample mean based on sampling iterations, with\n",
        "  warmup iterations discarded.\n",
        "  By default, we focus on the first parameter.\n",
        "  \"\"\"\n",
        "  # NOTE: why discard half of the samples?\n",
        "  cum_x = np.cumsum(x, axis)\n",
        "  return ((cum_x[1::2] - cum_x[:cum_x.shape[0]//2]) /\n",
        "         np.arange(1, cum_x.shape[0] // 2 + 1).reshape([-1] + [1] * (len(cum_x.shape) - 1)))\n",
        "\n",
        "long_error = mc_est(result_state_long.mean(1) - mean_est)\n",
        "short_error = mc_est(result_state_short.mean(1) - mean_est)\n",
        "\n",
        "true_var_available = True\n",
        "if (true_var_available):\n",
        "  target_precision = jnp.linalg.norm(var_est[0] / 100)\n",
        "else:\n",
        "  target_precision = jnp.linalg.norm(long_error[len(long_error) - 1], axis = -1)\n",
        "\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(jnp.linalg.norm(long_error, axis = -1), label = '4 chains')\n",
        "semilogy(jnp.linalg.norm(short_error, axis = -1), label = '1024 chains')\n",
        "hlines(target_precision, 0, total_samples_long / 2,\n",
        "       linestyles = '--', \n",
        "       label = 'Target: Var / 100')\n",
        "ylabel(\"Squared error for Mean estimate\")\n",
        "xlabel(\"Iterations (excluding warmup)\")\n",
        "legend(loc = 'best')\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOYqqpvA9iPE"
      },
      "source": [
        "I don't think the variance of the variance is stored in the inference gym, although it's probably possible to access this information using the error in the variance estimate. For now, we'll use the final result reported by the long chain as the target precision."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xDUsUDf66Hy"
      },
      "outputs": [],
      "source": [
        "long_var_error = mc_est(result_state_long.var(1)) - var_est\n",
        "short_var_error = mc_est(result_state_short.var(1)) - var_est\n",
        "\n",
        "long_var_estimate = jnp.linalg.norm(long_var_error[len(long_var_error) - 1], axis = -1)\n",
        "\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(jnp.linalg.norm(long_var_error, axis = -1), label = 'long')\n",
        "semilogy(jnp.linalg.norm(short_var_error, axis = -1), label = 'short')\n",
        "hlines(long_var_estimate, 0, total_samples_long / 2,\n",
        "       linestyles = '--', \n",
        "       label = 'long var estimate')\n",
        "ylabel(\"Squared error for Variance estimate\")\n",
        "legend(loc = 'best')\n",
        "show()\n",
        "\n",
        "# NOTE: why are the estimates in the long regime so poor??"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FF2FuZA-nIsm"
      },
      "source": [
        "Repeat the above, using a Monte Carlo estimator based on sampling iterations with only the warmup samples discard."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lTNQcUBiUXEF"
      },
      "outputs": [],
      "source": [
        "if (False):\n",
        "  print(result_state_long[num_warmup_long:, :, :].mean(0)[0][0])\n",
        "  print(result_state_short[num_warmup_short:, :, :].mean(0)[0][0])\n",
        "  print(mean_est[0])\n",
        "  print(long_error[len(long_error) - 1][0])\n",
        "  print(short_error[len(short_error) - 1][0])\n",
        "\n",
        "  result_state_long[num_warmup_long:, :, :].mean(1).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i21LwtfvqIVV"
      },
      "outputs": [],
      "source": [
        "def mc_est_warm(x, axis = 0):\n",
        "  \"\"\" compute running average without discarding half of the samples.\"\"\"\n",
        "  return np.cumsum(x, axis) / np.arange(1, x.shape[0] + 1).reshape([-1] + [1] * (len(x.shape) - 1))\n",
        "\n",
        "discard_warmup = True\n",
        "\n",
        "if (discard_warmup):\n",
        "  long_error = mc_est_warm(result_state_long[num_warmup_long:, :, :].mean(1)) - mean_est\n",
        "  short_error = mc_est_warm(result_state_short[num_warmup_short:, :, :].mean(1)) - mean_est\n",
        "else:\n",
        "  long_error = result_state_long[num_warmup_long:, :, :].mean(1) - mean_est\n",
        "  short_error = result_state_short[num_warmup_short:, :, :].mean(1) - mean_est\n",
        "\n",
        "true_var_available = True\n",
        "if (true_var_available):\n",
        "  target_precision = jnp.linalg.norm(var_est[0] / 100)\n",
        "else:\n",
        "  target_precision = jnp.linalg.norm(long_error[len(long_error) - 1], axis = -1)\n",
        "\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(jnp.linalg.norm(long_error, axis = -1), label = '4 chains')\n",
        "semilogy(jnp.linalg.norm(short_error, axis = -1), label = '512 chains')\n",
        "hlines(target_precision, 0, num_sampling_long,\n",
        "       linestyles = '--', \n",
        "       label = 'target: var / 100')\n",
        "ylabel(\"Squared error for Mean estimate\")\n",
        "xlabel(\"Sampling iterations (i.e. warmup excluded)\")\n",
        "legend(loc = 'best')\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iD3RgtL1-rNw"
      },
      "source": [
        "Remark: if after one iteration we are below the target precision, than we're probably running a warmup which is too long and / or running too many chains. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YSzP6JfpsVs-"
      },
      "outputs": [],
      "source": [
        "long_var_error = mc_est_warm(result_state_long[num_warmup_long:, :, :].var(1)) - var_est\n",
        "short_var_error = mc_est_warm(result_state_short[num_warmup_short:, :, :].var(1)) - var_est\n",
        "\n",
        "long_var_mc_estimate = jnp.linalg.norm(long_var_error[len(long_var_error) - 1], axis = -1)\n",
        "\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(jnp.linalg.norm(long_var_error, axis = -1), label = 'long')\n",
        "semilogy(jnp.linalg.norm(short_var_error, axis = -1), label = 'short')\n",
        "hlines(long_var_mc_estimate, 0, num_sampling_long,\n",
        "       linestyles = '--', \n",
        "       label = 'long MC estimate')\n",
        "ylabel(\"Squared error for Variance estimate\")\n",
        "legend(loc = 'best')\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DFL2yS4_ns01"
      },
      "source": [
        "Staring at the plot above it's clear that the short regime reaches a reasonable precision in fewer iterations than the long regime, even though the long regime warms up chains for many more iterations. The dotted line represent the Monte Carlo estimate using all the samples from the long regime. We'll use this as our target precision."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V-BRTgY8tUcr"
      },
      "outputs": [],
      "source": [
        "if (False):\n",
        "  print(long_mc_estimate)\n",
        "  print(jnp.linalg.norm(short_error, axis = -1)[0:10])\n",
        "  print(long_var_estimate)\n",
        "  print(jnp.linalg.norm(short_var_error, axis = -1)[0:10])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "enq2Z_mWoQvO"
      },
      "outputs": [],
      "source": [
        "# Identify the number of iterations after which the short regime matches\n",
        "# the precision of the long regime.\n",
        "# TODO: find a better criterion\n",
        "\n",
        "item_index = np.where(jnp.linalg.norm(short_error, axis = -1) \u003c= target_precision)\n",
        "target_iter_mean = item_index[0][0]\n",
        "print(\"Reasonable precision for mean reached in\", target_iter_mean + 1, \"iteration(s).\")\n",
        "\n",
        "item_index = np.where(jnp.linalg.norm(short_var_error, axis = -1) \u003c= long_var_estimate)\n",
        "target_iter_var = item_index[0][0]\n",
        "print(\"Reasonable precision for variance reached in\", target_iter_var + 1, \"iteration(s).\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QbQrYMHUGVdK"
      },
      "source": [
        "### Check for convergence"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f153fXSvGdqw"
      },
      "source": [
        "Let's first examine whether we're passed the transient bias regime (we should be since we're discarding the warmup phase)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4x_PF2UVGjdO"
      },
      "outputs": [],
      "source": [
        "# Plot last-sample estimarors\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(jnp.linalg.norm(result_state_long.mean(1) - var_est, axis=-1),\n",
        "         label='Long mean Error')\n",
        "semilogy(jnp.linalg.norm(result_state_short.mean(1) - mean_est, axis=-1),\n",
        "         label='Short Mean Error')\n",
        "hlines(jnp.sqrt(var_est.sum() / 100), 0, total_samples_long, label='Norm of Posterior Scales / 10')\n",
        "legend(loc='best')\n",
        "xlabel('Iteration')\n",
        "ylabel('Norm of Error of Estimate')\n",
        "title(target.name)\n",
        "xlim([0, 200])\n",
        "show()\n",
        "\n",
        "# NOTE: Note sure what's going on here."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zyfHfAHpulEc"
      },
      "source": [
        "Making due diligence, let's look at the samples returned by both methods, after discarding the warmup iterations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_1ITVYfchcai"
      },
      "outputs": [],
      "source": [
        "plot(result_long.all_states[num_warmup_short:, :, 0].flatten(), \n",
        "     result_long.all_states[num_warmup_short:, :, 1].flatten(), '.', alpha = 0.2)\n",
        "title('Long regime')\n",
        "show()\n",
        "plot(result_long.all_states[num_warmup_long:total_samples_long, :10, 1])\n",
        "show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VYYxWYr0l797"
      },
      "outputs": [],
      "source": [
        "# NOTE: (for Banana problem) With 4 samples after warmup we already samples spread\n",
        "# out accross the parameter space.\n",
        "num_samples_plot = 4  # target_iter_mean\n",
        "\n",
        "plot(result_short.all_states[num_warmup_short:num_samples_plot + num_warmup_short, :, 0].flatten(), \n",
        "     result_short.all_states[num_warmup_short:num_samples_plot + num_warmup_short, :, 1].flatten(), '.', alpha = 0.2)\n",
        "\n",
        "title('Short regime')\n",
        "show()\n",
        "plot(result_short.all_states[num_warmup_short:100 + num_warmup_short, [10, 20, 100, 500, 1000], 1])\n",
        "show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7OL1k6IQIBII"
      },
      "outputs": [],
      "source": [
        "# REMARK: the mixing for the banana problem is slow. This is obvious if we\n",
        "# only plot the first few samples of each chain.\n",
        "num_samples_plot = 4  # target_iter_mean\n",
        "plot(result_short.all_states[num_warmup_short:, :, 0].flatten(), \n",
        "     result_short.all_states[num_warmup_short:, :, 1].flatten(), '.', alpha = 0.2)\n",
        "title('Short regime')\n",
        "show()\n",
        "plot(result_short.all_states[:, [1, 200, 400, 600, 800, 1000], 1])\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B_uufKVPgCYH"
      },
      "source": [
        "Let's compute $\\hat R$ as a function of iteration and pay attention to how quickly $\\hat R$ goes to 1 in both regimes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fjgz5XhvgNZJ"
      },
      "outputs": [],
      "source": [
        "# NOTE: the warmup is not stored.\n",
        "# NOTE: compute rhat for the samples on the original space, since these are\n",
        "# the quantities of interest.\n",
        "def compute_rhat(result_state, num_samples, num_warmup = 0):\n",
        "  return tfp.mcmc.potential_scale_reduction(result_state[num_warmup:num_warmup + num_samples + 1],\n",
        "                                          independent_chain_ndims = 1).T\n",
        "\n",
        "# TODO: do this without a for loop\n",
        "# WARNING: this cell takes a minute to run\n",
        "# TODO: use a single variable num_sampling, instead of num_sampling_long and\n",
        "# num_sampling_var.\n",
        "rhat_long = np.array([])\n",
        "rhat_short = np.array([])\n",
        "range_iter = range(2, num_sampling_long, 10) # range(2, num_samples, 8)\n",
        "\n",
        "# NOTE: depending on the problem, it can be interesting to look at both.\n",
        "# However, to be consistent with earlier analysis, the warmup samples should\n",
        "# be discarded.\n",
        "discard_warmup = True\n",
        "\n",
        "for i in range_iter:\n",
        "  if (discard_warmup):\n",
        "    discard_long = num_warmup_long\n",
        "    discard_short = num_warmup_short\n",
        "  else:\n",
        "    discard_long = 0\n",
        "    discard_short = 0\n",
        "  rhat_long = np.append(rhat_long, \n",
        "                        compute_rhat(result_state_long, i, discard_long)[0, ])\n",
        "  rhat_short = np.append(rhat_short,\n",
        "                          compute_rhat(result_state_short, i, discard_short)[0, ])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iqyRPxi1r5im"
      },
      "source": [
        "Remark: the $\\hat R$ estimate can be quite noisy, especially when computed with a small number of samples. One manifestation of this is the fact that $\\hat R \u003c 1$. In the German credit score model, $\\hat R$ is as low as 0.6!! When this is the case, $\\hat R$ will typically be large for other parameters. Hence, inspecting many parameters (presumably all of interest) can safeguard us against crying \"victory\" too early.\n",
        "\n",
        "This type of noise can explain why the change in $\\hat R$ isn't always quite monotone, sometimes with an increase at first, and then the expected decrease."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bZ4psG2VqThM"
      },
      "outputs": [],
      "source": [
        "result_snip = result_state_long[num_warmup_long:num_warmup_long + 2]\n",
        "tfp.mcmc.potential_scale_reduction(result_snip, independent_chain_ndims = 1).T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8vsQTzT4CTIg"
      },
      "outputs": [],
      "source": [
        "# Plot result\n",
        "figure(figsize = [6, 6])\n",
        "semilogy(np.array(range_iter), rhat_long - 1, label = '4 chains')\n",
        "semilogy(np.array(range_iter), rhat_short - 1, label = '512 chains')\n",
        "legend(loc = 'best')\n",
        "xlim([0, 500])\n",
        "ylabel(\"Rhat - 1\")\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HIyVSEqSBR6c"
      },
      "source": [
        "(Banana example) As expected, $\\hat R$ decreases with the number of iterations per chain, although crucially not with the *total number of samples*! As one might suspect, the short regime produces a less noisy estimate of $\\hat R$. To be more precise, we expect $\\hat R$ to decrease with the *effective sample size per chain*. Since the long regime benefits from a longer warmup, the effective sample size per iteration should be better, although it might not make a difference in this example.\n",
        "\n",
        "Crucially, $\\hat R$ as a convergence diagnostic isn't sensitive to the fact we are running many chains (although the estimator does become less noisy...).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rX67xBppDRqj"
      },
      "outputs": [],
      "source": [
        "# Compare Rhat at the point where both methods have reached a comparable squared\n",
        "# error.\n",
        "# NOTE: not super reliable -- sometimes rhat is noisy and goes to 1 (or below)\n",
        "# before jumping back up...\n",
        "index = np.where(range_iter \u003e target_iter_mean)[0][0]\n",
        "print(\"Rhat for short regime after hitting target precision:\", rhat_short[index])\n",
        "print(\"Rhat for long regime after hitting target precision:\", rhat_long[len(rhat_long) - 1])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tDFKHqB3wIJn"
      },
      "source": [
        "**Proposition:** Concerned with how noisy $\\hat R$ might be, let's use a bootstrap scheme to get a standard deviation on the estimator. The short regime should be amiable to this, since we can resample chains. Unfortunately, if we sample with replacement, we underestimate the between chain variance, because some of the chains are identical. One idea is to randomly sample a subset of the chains without replacement and compute $\\hat R$.\n",
        "This will overestimate the uncertainty in our calculations, since we have reduced the sample size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZZ9zYF_3wFYP"
      },
      "outputs": [],
      "source": [
        "n_bootstrap_samples = 64  # 64\n",
        "rhat_estimates = np.array([])\n",
        "n_sampling_iter = max(range_iter[index], 2)  # max(target_iter_mean, 2) # range_iter[index]\n",
        "\n",
        "for i in range(1, n_bootstrap_samples):\n",
        "  choose_samples_randomly = True\n",
        "  if (choose_samples_randomly):\n",
        "    bootstrap_sample = np.random.choice(np.array(range(1, num_chains_short + 1)),\n",
        "                                        n_bootstrap_samples, replace = False)\n",
        "                                        # num_chains_short // 16, replace = False)\n",
        "  else:\n",
        "    bootstrap_sample = np.array(range(1 + (i - 1) * n_bootstrap_samples, i * n_bootstrap_samples))\n",
        "  # print(bootstrap_sample)\n",
        "  # print(result_state_short[:, bootstrap_sample, :].shape)\n",
        "  rhat_estimates = np.append(rhat_estimates,\n",
        "      compute_rhat(result_state_short[:, bootstrap_sample, :], n_sampling_iter, num_warmup_short)[0, ])\n",
        "\n",
        "print(\"Mean rhat (short) = \", rhat_estimates.mean(), \"+/-\", rhat_estimates.std())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FWjJ_J-cOdza"
      },
      "source": [
        "## Nested $\\hat R$\n",
        "\n",
        "To remedy the identified issue, we propose to pool chains together in the short regime, thereby building super-chains, and then checking that the super chains are mixing.\n",
        "\n",
        "We index each sample by $n$ the iteration, $m$ the chain, and $k$ the cluster of chains, and write $\\theta^{(n, m, k)}$. The within-chain variance is estimated by\n",
        "$$\n",
        "  s^2_{km} = \\frac{1}{N - 1} \\sum_{n = 1}^N \\left (\\theta^{(nmk)} - \\bar \\theta^{(.mk)} \\right)^2.\n",
        "$$\n",
        "\n",
        "Next the between-chain variance, or within super chain variance is\n",
        "\\begin{eqnarray}\n",
        "  s^2_{k.} \u0026 = \u0026 \\frac{1}{M - 1} \\sum_{m = 1}^M \\left (\\bar \\theta^{(.mk)} - \\bar \\theta^{(..k)} \\right)^2,\n",
        "\\end{eqnarray}\n",
        "and the total variance for a super chain is\n",
        "\\begin{eqnarray}\n",
        "  S^2_k \u0026 = \u0026 \\frac{1}{M - 1} \\sum_{m = 1}^M \\left (\\bar \\theta^{(.mk)} - \\bar \\theta^{(..k)} \\right)^2 + \\frac{1}{M (N - 1)} \\sum_{m = 1}^M \\sum_{n = 1}^N \\left (\\theta^{(nmk)} - \\bar \\theta^{(.mk)} \\right)^2 \\\\\n",
        "  \u0026 = \u0026 s^2_{k.} + \\frac{1}{M} \\sum_{m = 1}^M s^2_{km}\n",
        "\\end{eqnarray}\n",
        "Notice that this calculation accounts for the fact the super-chain is made up of multiple chains.\n",
        "Finally the within-super-chain variance is estimated as\n",
        "$$\n",
        "W = \\frac{1}{K} \\sum_{k = 1}^K S^2_k.\n",
        "$$\n",
        "\n",
        "Now it remains to compute the between super-chain variance\n",
        "$$\n",
        "B = \\frac{1}{K - 1} \\sum_{k = 1}^K \\left (\\bar \\theta^{(..k)} - \\bar \\theta^{(...)} \\right)^2,\n",
        "$$\n",
        "yielding an estimate of the posterior variance\n",
        "$$\n",
        "  \\widehat{\\mathrm{var}}^+(\\theta) = B + W,\n",
        "$$\n",
        "which very much looks like the posterior variance estimate used in the in the long regime, except that I've been a bit more consistent about making the estimator unbiased. We then compute\n",
        "\n",
        "$$\n",
        "  \\hat R = \\sqrt{\\frac{\\widehat{\\mathrm{var}}^+(\\theta)}{W}}.\n",
        "$$\n",
        "\n",
        "*Remark.* The $\\theta$ can be replaced by the rank-normalized $z$ as presrcribed by [Vehtari et al 2020](https://arxiv.org/abs/1903.08008).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FeuAkeGBQjln"
      },
      "source": [
        "\n",
        "### Implementation of nested-$\\hat R$ using TensorFlow.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s_6aGkHNHcjq"
      },
      "outputs": [],
      "source": [
        "# Remark: eager execution is disabled and would have to be enabled at the\n",
        "# start of the program. I however suspect this would interfere with\n",
        "# TensorFlow probability.\n",
        "tf.executing_eagerly()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jz0-no5z2LYo"
      },
      "outputs": [],
      "source": [
        "# Follow procedure described in source code for potential scale reduction.\n",
        "# NOTE: some of the tf argument need to be adjusted (e.g. keepdims = False,\n",
        "# instead of True). Not quite sure why.\n",
        "# QUESTION: can these be accessed as internal functions of tf?\n",
        "# TODO: following Pavel's example, rewrite this without using tf.\n",
        "# TODO: add error message when the number of samples is less than 2.\n",
        "\n",
        "# REMARK: this function doesn't seem to work, returns NaN.\n",
        "# As a result, can only use _reduce_variance with biased =  False.\n",
        "def _axis_size(x, axis = None):\n",
        "  \"\"\"Get number of elements of `x` in `axis`, as type `x.dtype`.\"\"\"\n",
        "  if axis is None:\n",
        "    return ps.cast(ps.size(x), x.dtype)\n",
        "  return ps.cast(\n",
        "      ps.reduce_prod(\n",
        "          ps.gather(ps.shape(x), axis)), x.dtype)\n",
        "\n",
        "def _reduce_variance(x, axis=None, biased=True, keepdims=False):\n",
        "  with tf.name_scope('reduce_variance'):\n",
        "    x = tf.convert_to_tensor(x, name='x')\n",
        "    mean = tf.reduce_mean(x, axis=axis, keepdims=True)\n",
        "    biased_var = tf.reduce_mean(\n",
        "        tf.math.squared_difference(x, mean), axis=axis, keepdims=keepdims)\n",
        "    if biased:\n",
        "      return biased_var\n",
        "    n = _axis_size(x, axis)\n",
        "    return (n / (n - 1.)) * biased_var\n",
        "\n",
        "def nested_rhat(result_state, num_super_chain):\n",
        "  used_samples = result_state.shape[0]\n",
        "  num_sub_chains = result_state.shape[1] // num_super_chains\n",
        "  num_dimensions = result_state.shape[2]\n",
        "\n",
        "  chain_states = result_state.reshape(used_samples, -1, num_sub_chains,\n",
        "                                      num_dimensions)\n",
        "\n",
        "  state = tf.convert_to_tensor(chain_states, name = 'state')\n",
        "  mean_chain = tf.reduce_mean(state, axis = 0)\n",
        "  mean_super_chain = tf.reduce_mean(state, axis = [0, 2])\n",
        "  variance_chain = _reduce_variance(state, axis = 0, biased = False)\n",
        "  variance_super_chain = _reduce_variance(mean_chain, axis = 1, biased = False) \\\n",
        "     + tf.reduce_mean(variance_chain, axis = 1)\n",
        "\n",
        "  W = tf.reduce_mean(variance_super_chain, axis = 0)\n",
        "  B = _reduce_variance(mean_super_chain, axis = 0, biased = False)\n",
        "\n",
        "  return tf.sqrt((W + B) / W)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PRDhSYB5SuG-"
      },
      "source": [
        "### CASE 1 (sanity check): $\\hat R$ after a few iterations\n",
        "The super chains are such that they have the same number of samples as the chains in the long regime. Because of the slow mixing, 4 iterations per chain is not enough to overcome the transient bias and the nested Rhat is high, even though each super chain has many iterations. Note we're looking at the first warmup iterations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EPUK6pcEZLSa"
      },
      "outputs": [],
      "source": [
        "# num_super_chains = 4\n",
        "# super_chain_size = num_chains_short // num_super_chains  # 250\n",
        "used_samples = 4  # total_samples_long // super_chain_size  # 4\n",
        "result_state = result_short.all_states[0:used_samples, :, :]\n",
        "\n",
        "print(\"short rhat: \", nested_rhat(result_state, num_super_chains))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VdhmxBJSS_q-"
      },
      "source": [
        "### CASE 2: $\\hat R$ after \"enough\" iterations\n",
        "The number of iterations in each chain corresponds to the number of samples required by the short regime to match the precision for the mean attained by the long regime after 1000 sampling iterations (meaning we've discarded the warmup iterations). The diagnostic is quite happy, even though there are only two iterations per chain."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F82uOj0gL8QM"
      },
      "outputs": [],
      "source": [
        "result_state.shape\n",
        "target_iter_mean"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QmUR7QNhiozI"
      },
      "outputs": [],
      "source": [
        "used_samples = max(target_iter_mean, 2)\n",
        "result_state = result_short.all_states[num_warmup_short:num_warmup_short + used_samples, :, :]\n",
        "\n",
        "print(\"short nested-rhat: \", nested_rhat(result_state, num_super_chains)[0])\n",
        "print(\"short rhat: \", rhat_short[index])\n",
        "print(\"long rhat: \", rhat_long[len(rhat_long) - 1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vx-f1_yL0OWb"
      },
      "outputs": [],
      "source": [
        "print(range_iter)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zrhot2Zc9g4T"
      },
      "outputs": [],
      "source": [
        "# Let's find out how quickly nested-rhat compared to traditional rhat goes down.\n",
        "nested_rhat_short = np.array([])\n",
        "for i in range_iter:\n",
        "    nested_rhat_short = np.append(nested_rhat_short, \n",
        "                                  nested_rhat(result_short.all_states[num_warmup_short:num_warmup_short + i, :, :],\n",
        "                                              num_super_chains).numpy()[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qAIWGYPhAV9E"
      },
      "outputs": [],
      "source": [
        "figure(figsize = [6, 6])\n",
        "semilogy(np.array(range_iter), rhat_long - 1, label = '$\\hat R$, 4 chains')\n",
        "semilogy(np.array(range_iter), rhat_short - 1, label = '$\\hat R$, 512 chains')\n",
        "semilogy(np.array(range_iter), nested_rhat_short - 1, label = '$n \\hat R$, 512 chains')\n",
        "legend(loc = 'best')\n",
        "xlim([0, 1000])\n",
        "ylabel(\"Rhat - 1\")\n",
        "xlabel(\"Post-warmup sampling iterations\")\n",
        "show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1S7mLM3fVaeA"
      },
      "outputs": [],
      "source": [
        "\n",
        "threshold = 1.1\n",
        "index_classic = np.where((rhat_short \u003c threshold) \u0026 (rhat_short \u003e 1.))\n",
        "if (len(index_classic[0]) \u003e 0):\n",
        "  print(\"Rhat =\", threshold, \"after\",range_iter[index_classic[0][0]], \"iterations.\")\n",
        "else:\n",
        "  print(\"Rhat doesn't hit the target threshold = \", threshold, \".\")\n",
        "\n",
        "\n",
        "index_short = np.where((nested_rhat_short \u003c threshold) \u0026 (nested_rhat_short \u003e 1.))\n",
        "if (len(index_short[0]) \u003e 0):\n",
        "  print(\"Nested Rhat =\", threshold, \"after\", range_iter[index_short[0][0]], \"iterations.\")\n",
        "else:\n",
        "  print(\"Nested Rhat doesn't hit the target threshold = \", threshold, \".\")\n",
        "\n",
        "threshold = 1.01\n",
        "index_classic = np.where((rhat_short \u003c threshold) \u0026 (rhat_short \u003e 1.))\n",
        "if (len(index_classic[0]) \u003e 0):\n",
        "  print(\"Rhat =\", threshold, \"after\",range_iter[index_classic[0][0]], \"iterations.\")\n",
        "else:\n",
        "  print(\"Rhat doesn't hit the target threshold = \", threshold, \".\")\n",
        "\n",
        "index_short = np.where((nested_rhat_short \u003c threshold) \u0026 (nested_rhat_short \u003e 1.))\n",
        "if (len(index_short[0]) \u003e 0):\n",
        "  print(\"Nested Rhat =\", threshold, \"after\", range_iter[index_short[0][0]], \"iterations.\")\n",
        "else:\n",
        "  print(\"Nested Rhat doesn't hit the target threshold = \", threshold, \".\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cLVLhU6Ydiia"
      },
      "source": [
        "# Effective sample size\n",
        "We'll now compute the effective sample size. We might in fact expect the classic diagnostic to work relatively well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2r8UQaRudkiH"
      },
      "outputs": [],
      "source": [
        "ess_long = np.sum(tfp.mcmc.effective_sample_size(\n",
        "    result_state_long[num_warmup_long:, : , :]), axis = 0)\n",
        "\n",
        "ess_short = np.sum(tfp.mcmc.effective_sample_size(\n",
        "    result_state_short[num_warmup_short:, :, :]), axis = 0)\n",
        "\n",
        "ess_short_target = np.sum(tfp.mcmc.effective_sample_size(\n",
        "    result_state_short[num_warmup_short:num_warmup_short + 3, :, :]), axis = 0)\n",
        "\n",
        "# NOTE: it seems we need at least 3 samples to compute the ess estimate...\n",
        "\n",
        "print(\"Ess long (discarding warmup): \", ess_long[0])\n",
        "print(\"Ess short (discarding warmup): \", ess_short[0])\n",
        "print(\"Ess short (when hitting target precision): \", ess_short_target[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fdt1EVb1nb5"
      },
      "source": [
        "# Adaptive warmup length\n",
        "Playing around a little, we find that once the algorithm is properly warmed up, the short regime can reach good precision in very few iterations. The primary limitation hence becomes the warm up time.\n",
        "\n",
        "Proper warmup means (i) we've overcomed the transient bias and have already moved across the \"typical set\" -- it isn't enough to be in the \"typical set\" if where we are is determined by our starting point -- and (ii) our algorithm tuned well-enough such that it can explore every part of the parameter space in a reasonable time and has a relatively short relaxation time. The first item is essential to both sampling regimes, though intuitively, it seems we might be able to compromise on the second item in the short regime.\n",
        "\n",
        "In many cases, the number of warmup samples is determined ahead of time when calling the algorithm. Ideally we'd stop the warmup once we have suitable tuning parameters and then move to the sampling phase. [Zhang et al (2020)](https://metrumrg.com/wp-content/uploads/2020/11/ZhangY_ACOP2020_BayesianWarmUpMethods.pdf) propose to run warmups over short windows of $w = 100$ iterations and compute $\\hat R$ and the ESS at the end of each of window to check if we should continue warming up. Once both diagnostic estimates are passed a certain threshold, the warmup ends and the sampling begins. In theory, this scheme can be adapted to the short regime by replacing $\\hat R$ with the nested $\\hat R$.\n",
        "\n",
        "My guess is that by using nested $\\hat R$ and the classic ESS (computed using many independent chains) we'll implicitly compromise on item (ii) -- so *a priori*, the described warmup method requires little adjustment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4v8M_CUWMSaQ"
      },
      "outputs": [],
      "source": [
        "# Define function to extract the adapted parameters\n",
        "# (Follow what's done in the inference gym tutorial)\n",
        "# REMARK: if we pass only initial step size, only one step size is adapted for \n",
        "# the whole transition kernel (as opposed to one step size per chain).\n",
        "# REMARK: we won't use this scheme. Instead, we'll pass the whole transition.\n",
        "from tensorflow_probability.python.internal.unnest import get_innermost\n",
        "\n",
        "\n",
        "# NOTE: presumable we're not going to use this, and instead get the full\n",
        "# kernel result back.\n",
        "def trace_fn(_, pkr):\n",
        "  return (\n",
        "      get_innermost(pkr, 'step_size'),\n",
        "      get_innermost(pkr, 'num_leapfrog_steps')\n",
        "      # get_innermost(pkr, 'max_trajectory_length')\n",
        "  )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ircvOe2_X0iY"
      },
      "outputs": [],
      "source": [
        "\n",
        "def forge_chain (target_rhat, warmup_window_size, kernel_cold, initial_state,\n",
        "                 max_num_steps, seed, monitor = False,\n",
        "                 use_nested_rhat = True, use_log_joint = False,\n",
        "                 num_super_chains = 4):\n",
        "  # store certain variables\n",
        "  rhat_forge = np.array([])\n",
        "  warmup_is_acceptable = False\n",
        "  store_results = []\n",
        "\n",
        "  warmup_iteration = 0\n",
        "\n",
        "  current_state = initial_state\n",
        "  final_kernel_args = None\n",
        "\n",
        "  while (not warmup_is_acceptable and warmup_iteration \u003c= max_num_steps):\n",
        "    warmup_iteration += 1\n",
        "\n",
        "    # 1) Run MCMC on short warmup window\n",
        "    result_cold, target_log_prob, final_kernel_args = tfp.mcmc.sample_chain(\n",
        "        num_results = warmup_window_size,\n",
        "        current_state = current_state,\n",
        "        kernel = kernel_cold,\n",
        "        previous_kernel_results = final_kernel_args,\n",
        "        seed = kernel_seed,\n",
        "        trace_fn = lambda _, pkr: unnest.get_innermost(pkr, 'target_log_prob'),\n",
        "        return_final_kernel_results = True)\n",
        "\n",
        "    if (warmup_iteration == 1) : \n",
        "      store_results = result_cold\n",
        "    else : \n",
        "      store_results = np.append(store_results, result_cold, axis = 0)\n",
        "\n",
        "    current_state = result_cold[-1]\n",
        "\n",
        "    # 2) Check if warmup is acceptable\n",
        "    if (used_nested_rhat):\n",
        "      if (use_log_joint):\n",
        "        shape_lp = target_log_prob.shape\n",
        "        rhat_warmup = nested_rhat(target_log_prob.reshape(shape_lp[0], shape_lp[1], 1),\n",
        "                                  num_super_chains)\n",
        "      else:\n",
        "        rhat_warmup = max(nested_rhat(result_cold, num_super_chains))\n",
        "    else:\n",
        "      if (use_log_joint):\n",
        "        rhat_warmup = tfp.mcmc.potential_scale_reduction(target_log_prob)\n",
        "      else:\n",
        "        rhat_warmup = max(tfp.mcmc.potential_scale_reduction(result_cold))\n",
        "    # ess_warmup = np.sum(tfp.mcmc.effective_sample_size(result_cold), axis = 0)\n",
        "\n",
        "    # print(rhat_warmup)\n",
        "\n",
        "    if (rhat_warmup \u003c target_rhat): warmup_is_acceptable = True\n",
        "    # if (max(rhat_warmup) \u003c 1.01 and min(ess_warmup) \u003e 100): warmup_is_acceptable = True\n",
        "\n",
        "    if (monitor):\n",
        "      print(\"step:\", final_kernel_args.step)\n",
        "      # print(\"max rhat:\", max(rhat_warmup))\n",
        "      # print(\"min ess warmup:\" , min(ess_warmup))\n",
        "      # print(\"step size:\", step_size)\n",
        "      # print(\"number of leapfrog steps:\", num_leapfrog_steps)\n",
        "  \n",
        "    save_values = True\n",
        "    if (save_values):\n",
        "      rhat_forge = np.append(rhat_forge, rhat_warmup)\n",
        "    # While loop ends\n",
        "\n",
        "  return store_results, final_kernel_args, rhat_forge\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zxmuUHohrdhj"
      },
      "outputs": [],
      "source": [
        "# Set up adaptive warmup scheme\n",
        "warmup_window_size = 5\n",
        "target_rhat = 1.01\n",
        "target_ess = 100\n",
        "max_num_steps = 1000 // warmup_window_size\n",
        "current_state = initial_state\n",
        "num_leapfrog_steps = 1\n",
        "warmup_iteration = 0\n",
        "kernel_seed = random.PRNGKey(1957)\n",
        "\n",
        "used_nested_rhat = True\n",
        "\n",
        "# define kernel using most recent step size\n",
        "kernel_cold = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "kernel_cold = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_cold, warmup_window_size)\n",
        "kernel_cold = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "      kernel_cold, warmup_window_size, target_accept_prob = 0.75,\n",
        "      reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "kernel_warm = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "kernel_warm = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_warm, 0)\n",
        "kernel_warm = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "      kernel_warm, warmup_window_size, target_accept_prob = 0.75,\n",
        "      reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZRrNCDrfdPHA"
      },
      "outputs": [],
      "source": [
        "result_cold, final_kernel_args, rhat_forge = \\\n",
        "  forge_chain(target_rhat = target_rhat,\n",
        "              warmup_window_size = warmup_window_size,\n",
        "              kernel_cold = kernel_cold,\n",
        "              initial_state = initial_state,\n",
        "              max_num_steps = max_num_steps,\n",
        "              seed = random.PRNGKey(1954), monitor = False,\n",
        "              use_nested_rhat = True,\n",
        "              use_log_joint = True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hbti3S9GOT6p"
      },
      "outputs": [],
      "source": [
        "print(\"iterations:\", len(rhat_forge) * warmup_window_size)\n",
        "print(rhat_forge)\n",
        "print(target_rhat)\n",
        "# print(tfp.mcmc.potential_scale_reduction(result_cold[-50]))\n",
        "# print(nested_rhat(result_short.all_states[num_warmup_short:num_warmup_short + 5, :, :], num_super_chains))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "14oysyyOvn3z"
      },
      "outputs": [],
      "source": [
        "# Run sampling iterations\n",
        "# def trace_fn(_, pkr):\n",
        "#   return (\n",
        "#       get_innermost(pkr, 'unnormalized_log_prob'))\n",
        "\n",
        "current_state = result_cold[-1]\n",
        "\n",
        "result_warm, target_log_prob, final_kernel_args_warm = tfp.mcmc.sample_chain(\n",
        "      num_results = 5,\n",
        "      current_state = current_state,\n",
        "      kernel = kernel_warm, # kernel_cold\n",
        "      previous_kernel_results = final_kernel_args,\n",
        "      seed = random.PRNGKey(100001),\n",
        "      return_final_kernel_results = True,\n",
        "      trace_fn = lambda _, pkr: unnest.get_innermost(pkr, 'target_log_prob'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nsN6o-g7tvQk"
      },
      "outputs": [],
      "source": [
        "print(tfp.mcmc.potential_scale_reduction(target_log_prob))\n",
        "# print(nested_rhat(target_log_prob, num_super_chains))\n",
        "shape_lp = target_log_prob.shape\n",
        "lp__ = target_log_prob.reshape(shape_lp[0], shape_lp[1], 1)\n",
        "lp__.shape\n",
        "print(nested_rhat(lp__, num_super_chains))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2BFNkJM-v7Oy"
      },
      "outputs": [],
      "source": [
        "print(tfp.mcmc.potential_scale_reduction(result_warm))\n",
        "nested_rhat(result_warm, num_super_chain = num_super_chains)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hcEJuAf2x9wt"
      },
      "outputs": [],
      "source": [
        "# options: result_cold[result_cold.shape[0] - 30:], result_state_short, result_warm, store_results\n",
        "states_to_read = result_warm\n",
        "\n",
        "print(\"mean estimate:\", np.mean(states_to_read.mean(0), axis = 0))\n",
        "print(\"variance estimate:\", np.mean(states_to_read.var(1), axis = 0))\n",
        "print(nested_rhat(states_to_read, num_super_chain = 4))\n",
        "print(tfp.mcmc.potential_scale_reduction(states_to_read))\n",
        "print(mean_est)\n",
        "print(var_est)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PehZnyZ6_S9y"
      },
      "outputs": [],
      "source": [
        "# Check output of the last run\n",
        "\n",
        "plot(result_warm[:, :, 0].flatten(), \n",
        "     result_warm[:, :, 1].flatten(), '.', alpha = 0.2)\n",
        "title('Long regime')\n",
        "show()\n",
        "plot(result_warm[:, :30, 1])\n",
        "show()\n",
        "\n",
        "# Compare to output we get with uninterrupted run.\n",
        "# (Examine the iterations before the warmup ends)\n",
        "chain_state_short = result_short.all_states[num_warmup_short - 10:num_warmup_short - 10 + warmup_window_size, :, :]\n",
        "plot(chain_state_short[:, :, 0].flatten(),\n",
        "     chain_state_short[:, :, 1].flatten(), '.', alpha = 0.2)\n",
        "show()\n",
        "\n",
        "plot(chain_state_short[:, :30, 1])\n",
        "show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PXfTAx-Qo6L0"
      },
      "source": [
        "## Experiment with window size"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0zDD4vvWDxuo"
      },
      "source": [
        "The code below returns the length of the warmup phase, simulated accross several seeds. This can give us a sense of how long the warmup phase is on average for different seeds. Be mindful that when using too many seeds with a lot of chains, the GPU can run out of memory. The motivation is to check how stable the warmup strategy is when using different window sizes.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bJOb0SpZo9BQ"
      },
      "outputs": [],
      "source": [
        "target_rhat = 1.01\n",
        "warmup_window_size = 30\n",
        "max_num_steps = 1000 // warmup_window_size\n",
        "\n",
        "iteration_after_warmup = np.array([])\n",
        "\n",
        "for seed in jax.random.split(jax.random.PRNGKey(0), 10):\n",
        "  initial_state = initialize((num_super_chains,), key = seed)\n",
        "  initial_state = np.repeat(initial_state, num_chains_short // num_super_chains,\n",
        "                            axis = 0)\n",
        "\n",
        "  result_cold, final_kernel_args, rhat_forge = \\\n",
        "    forge_chain(target_rhat = target_rhat,\n",
        "                warmup_window_size = warmup_window_size,\n",
        "                kernel_cold = kernel_cold,\n",
        "                initial_state = initial_state,\n",
        "                max_num_steps = max_num_steps,\n",
        "                seed = seed, monitor = False,\n",
        "                use_nested_rhat = True,\n",
        "                use_log_joint = False)\n",
        "    \n",
        "  iteration_after_warmup = np.append(iteration_after_warmup,\n",
        "                                     len(rhat_forge) * warmup_window_size)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ytyOSqQnFU0t"
      },
      "outputs": [],
      "source": [
        "# print(iteration_after_warmup)\n",
        "print(rhat_forge)\n",
        "print(iteration_after_warmup.mean())\n",
        "print(iteration_after_warmup.std())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AcJQmB6EreBE"
      },
      "source": [
        "### Results for the Banana problem\n",
        "Applying the code above for the banana problem with\n",
        "```\n",
        "target_rhat = 1.01\n",
        "use_nested_rhat = True\n",
        "use_log_joint = False\n",
        "```\n",
        "we estimate the length of the warmup phase for different window sizes:\n",
        "```\n",
        "w = 10, length = 62 +/- 16.12\n",
        "w = 15, length = 72 +/- 17.41\n",
        "w = 20, length = 86 +/- 18\n",
        "w = 30, length = 90 +/- 13.75\n",
        "w = 60, length = 120 +/- 0.0\n",
        "```\n",
        "Taking into consideration the different granularities, we find the results to be fairly consistent with one another.\n",
        "\n",
        "Let's go back to the original case where we use $\\hat R$ and ESS as our stoping criterion. Given the approximate one-to-one map between $\\hat R$ and ESS per chain, the two criterion are somewhat redundant, so I'll focus on $\\hat R$. When picking the window size, we must contend with the following trade-off:\n",
        "* if the window size is too short, we're unlikely to produce a large enough ESS per chain to hit the target $\\hat R$, and this could mean a never-ending warmup phase, or one that only stops once we exceed a maximum number of steps.\n",
        "* if the window size is too large, we may jump pass the optimal point. It's also worth noting that the first window is unlikely to yield satisfactory results, because the intial estimates are overdispered and bias. \n",
        "\n",
        "The first item is largely mitigated by using nested-$\\hat R$, since we're then less dependent on the ESS per chain. The second item could be addressed by using a path-finder to initialize the chains and/or by discarding some of the early iterations in a window when computing the diagnostics. \n",
        "\n",
        "One final remark is that using $\\hat R$ on the log joint distribution yielded somewhat optimistic results. As Pavel puts it: \"log_joint is a pretty bad metric. Generally, for convergence, you prefer to measure the least constrainted directions, and log_joint is typically not that.\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N11QsuHL8CsG"
      },
      "source": [
        "## Draft Code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WIYQh7xGVmdd"
      },
      "outputs": [],
      "source": [
        " \n",
        "result_cold, _, final_kernel_args = tfp.mcmc.sample_chain(\n",
        "      num_results = 100,\n",
        "      current_state = initial_state,\n",
        "      kernel = kernel_cold,\n",
        "      previous_kernel_results = None,\n",
        "      seed = random.PRNGKey(1954),\n",
        "      return_final_kernel_results = True)\n",
        "\n",
        "result_warm, _, final_kernel_args = tfp.mcmc.sample_chain(\n",
        "    num_results = 50,\n",
        "    current_state = result_cold[-1],\n",
        "    kernel = kernel_warm,\n",
        "    previous_kernel_results = final_kernel_args,\n",
        "    seed = random.PRNGKey(1954),\n",
        "    return_final_kernel_results = True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f-1qMw_GaPes"
      },
      "outputs": [],
      "source": [
        "nested_rhat(result_warm[1:3], 4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oqU_ci9mcZAP"
      },
      "outputs": [],
      "source": [
        "warmup_window_size = 200\n",
        "current_state = initial_state\n",
        "\n",
        "kernel_warm = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "kernel_warm = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_warm, warmup_window_size)\n",
        "kernel_warm = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "    kernel_warm, warmup_window_size, target_accept_prob = 0.75,\n",
        "    reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "# result_warm, (step_size_saved, num_leapfrog_steps_saved) = tfp.mcmc.sample_chain(\n",
        "#   warmup_window_size, current_state, kernel = kernel_warm,\n",
        "#   seed = random.PRNGKey(1954), trace_fn = trace_fn)\n",
        "\n",
        "result_warm, kernel_args, final_kernel_args = tfp.mcmc.sample_chain(\n",
        "    warmup_window_size, current_state, kernel = kernel_warm,\n",
        "    seed = random.PRNGKey(1954), return_final_kernel_results = True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RwXJqpjhc_NV"
      },
      "outputs": [],
      "source": [
        "# step_size = step_size_saved[warmup_window_size - 1]\n",
        "# current_state = result_warm[warmup_window_size - 1, :, :]\n",
        "# num_leapfrog_steps = num_leapfrog_steps_saved[warmup_window_size - 1]\n",
        "\n",
        "tfp.mcmc.potential_scale_reduction(result_warm[:, :, :])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bDEwArt3eJrx"
      },
      "outputs": [],
      "source": [
        "# kernel_warm2 = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, step_size, num_leapfrog_steps)\n",
        "# kernel_warm2 = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_warm2, warmup_window_size)\n",
        "# kernel_warm2 = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "#     kernel_warm2, warmup_window_size, target_accept_prob = 0.75,\n",
        "#     reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n",
        "\n",
        "# result_warm2, (step_size_saved) = tfp.mcmc.sample_chain(\n",
        "#   warmup_window_size, current_state, kernel = kernel_warm2,\n",
        "#   seed = random.PRNGKey(1954), trace_fn = trace_fn)\n",
        "\n",
        "result_warm2 = tfp.mcmc.sample_chain(\n",
        "    num_results = warmup_window_size, \n",
        "    kernel = kernel_warm,\n",
        "    current_state = current_state,\n",
        "    previous_kernel_results = final_kernel_args,\n",
        "    seed = random.PRNGKey(1953)\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KSpimtjyee-C"
      },
      "outputs": [],
      "source": [
        "tfp.mcmc.potential_scale_reduction(result_warm2.all_states[:, :, :])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Waq14PScSLdw"
      },
      "outputs": [],
      "source": [
        "print(problem_name)\n",
        "\n",
        "print(max(rhat_warmup))\n",
        "print(min(ess_warmup))\n",
        "# print(len(step_size))\n",
        "# print(step_size[0][warmup_window_size - 1])\n",
        "max(tfp.mcmc.potential_scale_reduction(result_warm))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9NKEcffA9oYJ"
      },
      "outputs": [],
      "source": [
        "# Define kernel for warmup windows (should be the same in the long and short regime)\n",
        "warmup_window_size = 10\n",
        "\n",
        "if (problem_name == 'Bananas' or problem_name == 'GermanCredit'):\n",
        "  kernel_warm_init = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, init_step_size, 1)\n",
        "  kernel_warm_init = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_warm_init, warmup_window_size)\n",
        "  kernel_warm_init = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "      kernel_warm_init, warmup_window_size, target_accept_prob = 0.75,  #0.75,\n",
        "      reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tpYvqf8jLH20"
      },
      "outputs": [],
      "source": [
        "result_warm, (step_size, max_trajectory_length) = tfp.mcmc.sample_chain(\n",
        "  warmup_window_size, initial_state, kernel = kernel_warm_init, seed = random.PRNGKey(1954),\n",
        "  trace_fn = trace_fn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J3WTFcu9FHoI"
      },
      "outputs": [],
      "source": [
        "print(step_size[len(step_size) - 1])\n",
        "print(max_trajectory_length[len(max_trajectory_length) - 1])\n",
        "print(initial_state.shape)\n",
        "print(result_warm[warmup_window_size, :, :])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UfB5JyVq_RYu"
      },
      "outputs": [],
      "source": [
        "# To run next window, define a new transition kernel\n",
        "# REMARK: the maximum trajectory length isn't, if my understanding is correct,\n",
        "# a tuning parameter; rather something that get's calculated at each step. So\n",
        "# there's no need to pass it on.\n",
        "kernel_warm = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn, step_size[len(step_size) - 1], 1)\n",
        "kernel_warm = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(kernel_warm, warmup_window_size)\n",
        "kernel_warm = tfp.mcmc.DualAveragingStepSizeAdaptation(\n",
        "    kernel_warm_init, warmup_window_size, target_accept_prob = 0.75,\n",
        "    reduce_fn = tfp.math.reduce_log_harmonic_mean_exp)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mtoniiOuH96h"
      },
      "outputs": [],
      "source": [
        "result_warm2, (step_size, max_trajectory_length) = tfp.mcmc.sample_chain(\n",
        "  warmup_window_size, initial_state, kernel = kernel_warm, seed = random.PRNGKey(1954),\n",
        "  trace_fn = trace_fn)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n2srI_MGBpCD"
      },
      "outputs": [],
      "source": [
        "print(result_warm.shape)\n",
        "print(step_size.shape)\n",
        "print(max_trajectory_length.shape) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qfP3IRJwAvEG"
      },
      "outputs": [],
      "source": [
        "step_size[len(step_size) - 1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DYx5bInSJw3o"
      },
      "outputs": [],
      "source": [
        "# nested_rhat(result_short.all_states, num_super_chains)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s2GyNaFntSnv"
      },
      "outputs": [],
      "source": [
        "## Sandbox\n",
        "\n",
        "# Pool chains into super chains\n",
        "# num_super_chains = 4  # num_chains_short // num_chains_long\n",
        "# num_sub_chains = num_chains_short // num_super_chains\n",
        "# used_samples = num_samples # 5  # 2 * target_iter_mean  # target_iter_mean\n",
        "# result_state = result_short.all_states[0:used_samples, :, :]\n",
        "# chain_states = result_state.reshape(used_samples, num_sub_chains,\n",
        "#                                     -1, num_dimensions)\n",
        "\n",
        "# independent_chains_ndims = 1\n",
        "# sample_ndims = 1\n",
        "# sample_axis = tf.range(0, sample_ndims)\n",
        "# chain_axis "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LSrGbmX7knkn"
      },
      "outputs": [],
      "source": [
        "  # used_samples = result_state.shape[0]\n",
        "  # num_sub_chains = result_state.shape[1] // num_super_chains\n",
        "  # num_dimensions = result_state.shape[2]\n",
        "\n",
        "  # chain_states = result_state.reshape(used_samples, -1, num_sub_chains,\n",
        "  #                                    num_dimensions)\n",
        "\n",
        "  # state = tf.convert_to_tensor(chain_states, name = 'state')\n",
        "\n",
        "  # mean_chain = tf.reduce_mean(state, axis = 0)\n",
        "  # mean_super_chain = tf.reduce_mean(state, axis = [0, 2])\n",
        "  # variance_chain = _reduce_variance(state, axis = 0, biased = False)\n",
        "  # variance_super_chain = _reduce_variance(mean_chain, axis = 1, biased = False) \\\n",
        "  #    + tf.reduce_mean(variance_chain, axis = 1)\n",
        "\n",
        "  # W = tf.reduce_mean(variance_super_chain, axis = 0)\n",
        "  # B = _reduce_variance(mean_super_chain, axis = 0, biased = False)\n",
        "\n",
        "  # rhat = tf.sqrt((W + B) / W)\n",
        "\n",
        "  # print(rhat)\n",
        "\n",
        "  # print(mean_chain.shape)\n",
        "  # print(mean_super_chain.shape)\n",
        "  # print(\"mean_super_chain: \", mean_super_chain)\n",
        "  # print(variance_chain.shape)\n",
        "  # print(variance_super_chain.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G2ktNw3wpvXv"
      },
      "outputs": [],
      "source": [
        "# print(state.shape)  # (5, 250, 4, 2)\n",
        "# print(result_state.shape)  # (5, 1000, 2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P9cllfr7k10F"
      },
      "outputs": [],
      "source": [
        "# # 'manually' compute the mean of each super chain.\n",
        "# print(np.mean(result_state[:, 0:250, 0]))\n",
        "# print(np.mean(result_state[:, 250:500, 0]))\n",
        "# print(np.mean(result_state[:, 500:750, 0]))\n",
        "# print(np.mean(result_state[:, 750:1000, 0]))\n",
        "\n",
        "# # compute the means after reshaping the results. Get agreement!\n",
        "# print(np.mean(chain_states[:, 0, :, 0]))\n",
        "# print(np.mean(chain_states[:, 1, :, 0]))\n",
        "# print(np.mean(chain_states[:, 2, :, 0]))\n",
        "# print(np.mean(chain_states[:, 3, :, 0]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UmAwXIlY0KPE"
      },
      "outputs": [],
      "source": [
        "# print(result_state[:, 250, 0])\n",
        "# print(chain_states[:, 0, 1, 0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UrFyfk4Ms4iJ"
      },
      "outputs": [],
      "source": [
        "# simple_chain = np.array([[0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4, 5]])\n",
        "# simple_chain.shape  # (4, 6)\n",
        "\n",
        "# chain_reshape = simple_chain.reshape(4, 2, -1)\n",
        "# chain_reshape.shape  # (4, 2, 3)\n",
        "# np.mean(chain_reshape, axis = 0)  # returns mean for each chain\n",
        "# np.mean(chain_reshape[:, 0, :])  # 1\n",
        "# np.mean(chain_reshape[:, 1, :])  # 4\n",
        "\n",
        "# np.mean(simple_chain[:, 0:2])  # 1\n",
        "# np.mean(simple_chain[:, 3:6])  # 4 -- but it seems index should be 3:5\n",
        "# # simple_chain[:, 3:6]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VhdfCADtCKMp"
      },
      "outputs": [],
      "source": [
        "# ## Sandbox\n",
        "\n",
        "# tf.compat.v1.disable_eager_execution() # need to disable eager in TF2.x\n",
        "\n",
        "# state = result_short.all_states[1:range_iter[index], :, :]\n",
        "# n = state.shape[0]\n",
        "# m = state.shape[1]\n",
        "\n",
        "# sample_ndims = 1\n",
        "# independent_chains_ndims = 1\n",
        "# sample_axis = tf.range(0, sample_ndims)  # CHECK\n",
        "# chain_axis = 0\n",
        "# sample_and_chain_axis = tf.range(0, sample_ndims + independent_chains_ndims)  # CHECK\n",
        "\n",
        "\n",
        "# with tf.name_scope('potential_scale_reduction_single_state'):\n",
        "#   state = tf.convert_to_tensor(state, name = 'state')\n",
        "\n",
        "#   # CHECK: do we need to define a tf scope?\n",
        "#   n_samples = tf.compat.dimension_value(state.shape[0])\n",
        "\n",
        "#   # n = _axis_size(state, sample_axis)\n",
        "#   # m = _axis_size(state, chain_axis)\n",
        "\n",
        "#   # NOTE: These lines prompt the error message once the session is run.\n",
        "#   # x = tf.reduce_mean(state, axis=sample_axis, keepdims=True)\n",
        "#   # x_tf = tf.convert_to_tensor(x, name = 'x')\n",
        "#   # n_tf = _axis_size(x_tf)\n",
        "\n",
        "#   b_div_n = _reduce_variance(\n",
        "#       tf.reduce_mean(state, axis = 0, keepdims = False),\n",
        "#       sample_and_chain_axis,  # sample and chain axis\n",
        "#       biased = False\n",
        "#   )\n",
        "\n",
        "#   w = tf.reduce_mean(\n",
        "#       _reduce_variance(state, sample_axis, keepdims = False, \n",
        "#                        biased = False),\n",
        "#       axis = sample_and_chain_axis\n",
        "#   )\n",
        "\n",
        "#   # TODO: work out n and m from the number of chains being passed.\n",
        "#   # n = target_iter_mean\n",
        "#   # m = num_chains\n",
        "#   sigma_2_plus = ((n - 1) / n) * w + b_div_n\n",
        "#   rhat = ((m + 1.) / m) * sigma_2_plus / w - (n - 1.) / (m * n)\n",
        "\n",
        "\n",
        "# # Launch the graph in a session. (TensorFlow uses differed action,\n",
        "# # so need to explicitly request evaluation)\n",
        "# sess = tf.compat.v1.Session()\n",
        "\n",
        "# print(sess.run(rhat))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "rhat_locker.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "17lXD994FFhRE3RcfFaTCTyRyzF0cXvkP",
          "timestamp": 1632344011320
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
